#include <float.h> // 引入float.h以使用浮点数相关的常量，如FLT_MAX

// CUDA核函数，用于执行2D最大池化操作
// 参数说明：
// dev_input: 输入的设备端（GPU）浮点型数组，表示输入的图像数据
// dev_out: 输出的设备端（GPU）浮点型数组，表示最大池化后的结果
/*
    N: 批次大小，表示有多少个样本，例如，在深度学习中，输入的数据通常会被组织成多个批次来进行并行处理。每个批次包含多个样本（图片）。
    N 就表示每个批次中有多少个样本。
*/

/* 
    C: 通道数，通常指图像的通道，如RGB的3个通道，C 就是指输入数据中每个样本（图像）有多少个通道。
    每个通道的最大池化操作是独立进行的，即每个通道会在自己的池化窗口内找最大值，而不会和其他通道的值混合在一起。
*/
// H: 输入图像的高度
// W: 输入图像的宽度
// H_o: 输出图像的高度（池化后的高度）
// W_o: 输出图像的宽度（池化后的宽度）
__global__ void maxpool2d(const float *dev_input, float *dev_out, int N, int C, int H, int W, int H_o, int W_o)
{
    // 计算当前线程的全局索引
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    // 输出的总元素数（所有批次、通道、输出图像的每个像素点）
    int total_outputs = N * C * H_o * W_o;

    // 如果线程索引超出了总输出范围，则直接返回（即该线程不做任何操作）
    if (idx >= total_outputs)
        return;

    // 通过全局索引计算当前线程负责的输出位置（w_o, h_o, c, n）
    // w_o: 输出图像的宽度索引
    int w_o = idx % W_o;
    // idx1: 去除宽度后的剩余索引部分
    int idx1 = idx / W_o;
    // h_o: 输出图像的高度索引
    int h_o = idx1 % H_o;
    // idx2: 去除高度后的剩余索引部分
    int idx2 = idx1 / H_o;
    // c: 通道索引
    int c = idx2 % C;
    // n: 批次索引
    int n = idx2 / C;

    // 计算输入图像中对应的起始高度和宽度
    // 假设池化窗口为2x2，因此乘以2表示原图对应的输入区域的起始坐标
    int h_in = h_o * 2;
    int w_in = w_o * 2;

    // 初始化最大值为最小浮点数，作为初始比较值
    float max_val = -FLT_MAX;

    // 双重循环，用于遍历池化窗口2x2内的元素
    // kh: 池化窗口的行偏移
    for (int kh = 0; kh < 2; ++kh)
    {
        // kw: 池化窗口的列偏移
        for (int kw = 0; kw < 2; ++kw)
        {
            // 计算当前池化窗口中元素的实际位置（在输入图像中的位置）
            int h = h_in + kh;
            int w = w_in + kw;

            // 检查输入的坐标是否超出输入图像的范围（边界处理）
            if (h < H && w < W)
            {
                // 根据批次n、通道c、高度h和宽度w计算输入图像中的索引
                int input_idx = ((n * C + c) * H + h) * W + w;
                // 获取当前输入图像中的值
                float val = dev_input[input_idx];

                // 如果当前值大于当前最大值，则更新最大值
                if (val > max_val)
                    max_val = val;
            }
        }
    }

    // 计算输出图像中的索引位置
    int output_idx = ((n * C + c) * H_o + h_o) * W_o + w_o;

    // 将最大值存储到输出数组的对应位置
    dev_out[output_idx] = max_val;
}
